#!/usr/bin/python

import sys
import os
import os.path
import re
from mailbox import mbox, Maildir, mboxMessage
import hashlib
from collections import defaultdict
from optparse import OptionParser

from parse_date import parse_date, get_year
from cleanup_froms_in_mbox import cleanup_mbox

class MailSorter:
    RESTRUCT = "restruct"
    FILEVALID = "file_valid"
    REMOVE_DUP = "remove_duplicates"
    
    available_modes = [FILEVALID, REMOVE_DUP, RESTRUCT]
    
    date_searcher = re.compile("(\n|^)Date:[^\n]*\n")
    
    def __init__(self, input_dirs, modes=None, out_fn=None, log_fn=None, err_fn=None,
                 odir=None, trash=None):
        self.mail_hash_db = defaultdict(list)
        self.input_dirs = input_dirs
        self.maildirs = set()
        self.tomove = set()
        
        if modes is not None:
            self.modes = modes
        else:
            self.modes = available_modes[:1]
        
        if out_fn is not None:
            self.outf = open(out_fn, "w")
        else:
            self.outf = sys.stdout
        
        if log_fn is not None:
            self.logf = open(log_fn, "w")
        else:
            self.logf = open(os.devnull, "w")
        
        if err_fn is not None:
            self.errf = open(log_fn, "w")
        else:
            self.errf = sys.stderr
        
        self.trashbox = trash
        
        if odir is not None:
            self.odir = odir

    @classmethod
    def hash(cls, a):
        return hashlib.sha512(a).hexdigest()
    
    @classmethod
    def payload_to_str(cls, msg):
        if msg.is_multipart():
            return " ".join((MailSorter.payload_to_str(m) for m in msg.get_payload()))
        else:
            return msg.get_payload()

    @classmethod
    def choose_best_message(cls, l):
        first = l[0]
        for msg in l:
            mbox_name, mbox_type, msg_date = msg
            if first[0].startswith("trash"):
                first = msg
                continue
            if msg_date is None:
                continue
            if msg_date < first[2]:
                first = msg
                continue
        return first
    
    @classmethod
    def mark_duplicates(cls, mail_list, h):
        d = defaultdict(set)
        if len(mail_list) == 1:
            return d
        bm = MailSorter.choose_best_message(mail_list)
        bm_name, bm_type, bm_date = bm
        best_index = mail_list.index(bm)
        for i, other in enumerate(mail_list):
            name, mboxtype, date = other
            if i == best_index:
                continue
            d[(name, mboxtype)].add((h, date, (name == bm_name) and (date == bm_date)))
        return d
    
    @classmethod
    def guess_date_line(cls, msg):
        s = cls.payload_to_str(msg)
        found = cls.date_searcher.search(s)
        if found:
            return found.group(0).strip()
        else:
            return
    
    @classmethod
    def guess_year_by_from(cls, msg):
        if isinstance(msg, mboxMessage):
            fr = msg.get_from()
            year = get_year(fr)
            return year
    
    def set_modes(self, modes):
        self.modes = modes
    
    def add_mails_to_hashdb(self, mbox):
        mbox, mbox_path, mbox_type = mbox
        for key in mbox.keys():
            msg = mbox[key]
            h = MailSorter.hash(MailSorter.payload_to_str(msg))
            self.mail_hash_db[h].append( (mbox_path, mbox_type, parse_date(msg["Date"])) )
    
    def add_mails_to_yeardb(self, mbox):
        mbox, mbox_name, src_path, src_type = mbox
        self.logf.write("mbox {0} has {1} mails\n".format(mbox_name, len(mbox.keys())))
        for key in mbox.keys():
            msg = mbox[key]
            if msg["Date"] is not None:
                date_line = msg["Date"]
            else:
                date_line = MailSorter.guess_date_line(msg)
            
            
            date = parse_date(date_line)
            if not date:
                year = get_year(date_line)
            else:
                year = date[0]
            
            if year == "Unsorted":
                from_year = MailSorter.guess_year_by_from(msg)
                if from_year is not None:
                    year = from_year
            
            h = MailSorter.hash(MailSorter.payload_to_str(msg))
            
            if not mbox_name in self.mail_year_db[year].keys():
                self.mail_year_db[year][mbox_name] = defaultdict(set)
            self.mail_year_db[year][mbox_name][(src_path, src_type)].add(h)
    
    def mbox_handler(self, filename, path):
        try:
            mybox = mbox(path)

            if len(mybox.keys()) == 0:
                self.logf.write("Nomail: " + path + "\n")
                self.tomove.add(path)
                return

            cleanup_mbox(mybox, self.logf)
            if MailSorter.REMOVE_DUP in self.modes:
                self.add_mails_to_hashdb((mybox, path, "mbox"))
            if MailSorter.RESTRUCT in self.modes:
                self.add_mails_to_yeardb((mybox, filename, path, "mbox"))
        except IOError, ioe:
            self.errf.write("IOError: " + str(ioe) + "\n")
            self.tomove.add(path)
            return
    
    def maildir_handler(self, d, cd):
        my_maildir = Maildir(d, factory=None)
        mybox_name = d.rstrip("/")
        if mybox_name.find("/") >= 0:
            mybox_name = d.strip("/").rsplit("/", 1)[1]
        
        #creating empty directories
        if "cur" not in cd:
            added_cur = os.path.join(d, "cur")
            os.mkdir(added_cur)
                
        if "new" not in cd:
            added_new = os.path.join(d, "new")
            os.mkdir(added_new)
            
        for a in ["new", "cur", "tmp"]:
            self.maildirs.add(os.path.join(d, a))
            
            if MailSorter.REMOVE_DUP in self.modes:
                self.add_mails_to_hashdb((my_maildir, d, "maildir"))
            if MailSorter.RESTRUCT in self.modes:
                self.add_mails_to_yeardb((my_maildir, mybox_name, d, "maildir"))
    
    def collect_info(self, idir):
        self.logf.write("Primary Working dir: {0}\n".format(idir))
        for d, cd, files in os.walk(idir):
            self.logf.write("Working dir: " + d + "\n")
            self.logf.flush()
            
            #under a Maildir directory
            is_maildir = False
            for md in self.maildirs:
                if d.startswith(md):
                    self.logf.write("Skipping Maildir subfolders in " + d + "\n")
                    self.logf.flush()
                    is_maildir = True
                    break
            if is_maildir:
                continue
            
            #mbox
            for f in files:
                p = os.path.join(d, f)
                self.mbox_handler(f, p)
            
            #Maildir
            if 'new' in cd or 'cur' in cd:
                self.logf.write("Maildir in " + d + "\n")
                self.logf.flush()
                self.maildir_handler(d, cd)
    
    
    def get_duplicates(self):
        what_to_remove = defaultdict(set)
        
        total = len(self.mail_hash_db)
        for count, h in enumerate(self.mail_hash_db):
            if 100 * count / total > 100 * (count - 1) / total:
                self.logf.write("marking duplicates: %d %%\n" % (100 * count / total))
                self.logf.flush()
            
            mail_list = self.mail_hash_db[h]
            my_dup = MailSorter.mark_duplicates(mail_list, h)
            for k in my_dup.keys():
                for item in my_dup[k]:
                    what_to_remove[k].add(item)
        return what_to_remove
    
    def _remove_duplicates(self, what_to_remove):
        if self.trashbox is not None:
            self.trashbox = mbox(self.trashbox)
        
        total = len(what_to_remove.keys())
        for count, (mybox_path, mybox_type) in enumerate(what_to_remove.keys()):
            
            if 100 * count / total > 100 * (count - 1) / total:
                self.logf.write("removing duplicates: {0} %%\n".format(
                    100 * count / total))
                self.logf.flush()
            
            mybox = None
            if mybox_type == "mbox":
                mybox = mbox(mybox_path)
            elif mybox_type == "maildir":
                mybox = Maildir(mybox_path, factory=None)
            
            done = set()
            notfirst = defaultdict(bool)
            for key in mybox.keys():
                mail = mybox[key]
                h = MailSorter.hash(MailSorter.payload_to_str(mail))
                date = parse_date(mail["Date"])
                
                if (h, date, True) in what_to_remove[(mybox_path, mybox_type)]:
                    if not notfirst[h]:
                        notfirst[h] = True
                        continue
                    if self.trashbox is not None:
                        self.trashbox.add(mail)
                    mybox.remove(key)
                    done.add((h, date))
                    self.logf.write("One mail removed.\n")
                    self.logf.flush()
                elif (h, date, False) in what_to_remove[(mybox_path, mybox_type)]:
                    if self.trashbox is not None:
                        self.trashbox.add(mail)
                    mybox.remove(key)
                    done.add((h, date))
                    self.logf.write("One mail removed.\n")
                    self.logf.flush()
               
            mybox.flush()
            mybox.close()
    
    def remove_duplicates(self):
        what_to_remove = self.get_duplicates()
        self._remove_duplicates(what_to_remove)
    
    def move_no_mail(self):
        from os import rename
        from os.path import basename, exists, isfile, isdir
        nmdir = os.path.join(self.odir, "NoMail")
        try:
            os.mkdir(nmdir)
        except OSError:
            self.logf.write("Writing into directory that already exists:")
            self.logf.write("{0}\n".format(nmdir))
            self.logf.flush()
        for p in self.tomove:
            name = basename(p.rstrip("/"))
            new = os.path.join(nmdir, name)
            if isdir(p):
                raise Exception("nomail file is a directory?!")
            if exists(new):
                if isfile(new):
                    hs = MailSorter.hash(file(p).read())
                    ht = MailSorter.hash(file(new).read())
                    if hs == ht:
                        os.remove(p)
                    else:
                        for i in map(str, range(10)) + map(chr, range(97, 123)):
                            new_filename = new + "-%s" % i
                            if os.path.exists(new_filename):
                                ht2 = MailSorter.hash(file(new_filename).read())
                                if ht2 == hs:
                                    os.remove(p)
                                    break
                            else:
                                os.rename(p, new_filename)
                                break
                else:
                    raise Exception("directory inside nomail dierctory")
            else:
                os.rename(p, new)


    def restruct(self):
        try:
            os.mkdir(self.odir)
        except OSError:
            self.logf.write("Writing into directory that already exists:")
            self.logf.write("{0}\n".format(self.odir))
            self.logf.flush()
        
        self.logf.write("There will be {0} year dirs.\n".format(
            len(self.mail_year_db.keys())))
        self.logf.flush()
        for year in self.mail_year_db.keys():
            self.logf.write("Writing {0}.\n".format(str(year)))
            self.logf.flush()
            try:
                os.mkdir(os.path.join(self.odir, str(year)))
            except OSError:
                self.logf.write("Writing into directory that already exists:")
                self.logf.write("{0}\n".format(os.path.join(self.odir, str(year))))
                self.logf.flush()
            for tag in self.mail_year_db[year].keys():
                outbox = mbox(os.path.join(self.odir, str(year), tag.replace("\\", "")))
                for mbox_data in self.mail_year_db[year][tag]:
                    mboxpath, mbox_type = mbox_data
                    if mbox_type == "mbox":
                        mybox = mbox(mboxpath)
                    elif mbox_type == "maildir":
                        mybox = Maildir(mboxpath, factory=None)
                    
                    for key in mybox.keys():
                        mail = mybox[key]
                        h = MailSorter.hash(MailSorter.payload_to_str(mail))
                        if h in self.mail_year_db[year][tag][mbox_data]:
                            outbox.add(mail)
                outbox.flush()
                outbox.close()
        if len(self.tomove) > 0:
            self.move_no_mail()
    
    def run(self):
        if MailSorter.RESTRUCT in self.modes:
            self.mail_year_db = defaultdict(dict)
        
        self.logf.write("Collecting mail data started.\n")
        self.logf.flush()
        for idir in self.input_dirs:
            self.collect_info(idir)
        self.logf.write("Collecting ended.\n")
        self.logf.flush()
        if MailSorter.REMOVE_DUP in self.modes:
            self.logf.write("Removing duplicates started.\n")
            self.logf.flush()
            self.remove_duplicates()
            self.logf.write("Removing duplicates ended.\n")
            self.logf.flush()
        if MailSorter.RESTRUCT in self.modes:
            self.logf.write("Restructing started.\n")
            self.logf.flush()
            self.restruct()
            self.logf.write("Restructing ended.\n")
            self.logf.flush()

def option_list_callback(option, opt, value, parser):
    if value is not None:
        setattr(parser.values, option.dest, [v.strip() for v in value.split(',')])


def create_parser():
    parser = OptionParser()


    parser.add_option("-m", "--modes", dest="mode", action="callback", type="string",
                      help="mode (file_valid/remove_duplicates/restruct)",
                      callback=option_list_callback)
    parser.add_option("-i", "--input_dir", dest="idir", action="callback", type="string",
                      help="input root directory", callback=option_list_callback)
    parser.add_option("-t", "--trash", dest="trash",
                      help="trash file name")
    parser.add_option("-o", "--output_dir", dest="odir",
                      help="output mbox root directory")
    parser.add_option("-e", "--errorfile", dest="errf",
                      help="file name for error messages (instead of stderr)")
    parser.add_option("-l", "--logfile", dest="logf",
                      help="file name for log messages")
    parser.add_option("", "--outfile", dest="outf",
                      help="file name for output messages (instead of stdout)")
    return parser

def main():
    parser = create_parser()
    (options, args) = parser.parse_args()

    if options.mode is None:
        parser.print_help()
        quit()

    available_modes = set(["file_valid", "remove_duplicates", "restruct"])
    if len(set(options.mode) & available_modes) == 0:
        sys.stderr.write("Not available mode. Options: %s\n" % (str(available_modes)))
        sys.exit(-1)

    if not options.idir:
        sys.stderr.write("Input directory not defined\n")
        sys.exit(-1)

    if options.mode == "restruct":
        if not options.odir:
            sys.stderr.write("No output directory given")
            sys.exit(-1)

    modes = options.mode
    out_fn = options.outf
    log_fn = options.logf
    err_fn = options.errf
    odir = options.odir
    trash = options.trash
    ms = MailSorter(options.idir, modes, out_fn=out_fn, log_fn=log_fn, err_fn=err_fn,
                     odir=odir, trash=trash)
    ms.run()

if __name__ == "__main__":
    main()

